#include "stereo_api.h"
#include <iostream>

using namespace cv;
using namespace std;

enum {
	STEREO_BM=0,
	STEREO_SGBM=1,
	STEREO_HH=2,
	STEREO_VAR=3,
	STEREO_3WAY=4
};


//use stereo parameters to correct images
void correctStereoImages(const Mat &leftGray, const Mat &rightGray,
			Mat &cameraMatrixLeft, Mat &cameraMatrixRight, Mat &discoeffLeft, Mat &discoeffRight, Mat &R, Mat &T,
			Mat &rectifyImageL, Mat &rectifyImageR, Rect &validRoil, Rect &validRoir, Mat &q)
{
	Mat rL, rR, pL, pR;
	Mat mapLx, mapLy, mapRx, mapRy;

	Size imageSize = leftGray.size();
	stereoRectify(cameraMatrixLeft, discoeffLeft, cameraMatrixRight, discoeffRight, imageSize,
						R, T, rL, rR, pL, pR, q, CALIB_ZERO_DISPARITY, -1, imageSize, &validRoil, &validRoir);

	initUndistortRectifyMap(cameraMatrixLeft, discoeffLeft, rL, pL, imageSize, CV_32FC1, mapLx, mapLy);
	initUndistortRectifyMap(cameraMatrixRight, discoeffRight, rR, pR, imageSize, CV_32FC1, mapRx, mapRy);

	remap(leftGray, rectifyImageL, mapLx, mapLy, INTER_LINEAR);
	remap(rightGray, rectifyImageR, mapRx, mapRy, INTER_LINEAR);
}

void stereoMatch(Rect &validRoil, Rect &validRoir, Mat &rectifyImageL, Mat &rectifyImageR, Mat &leftDisp)
{
	int alg = STEREO_SGBM;
	int SADWindowSize, numberOfDisparities;
	SADWindowSize = 9;
	numberOfDisparities = ((rectifyImageL.size().width/8) + 15) & -16;

	Ptr<StereoBM> bm = StereoBM::create(16,9);
	Ptr<StereoSGBM> sgbm = StereoSGBM::create(0,16,3);

	//set bm
	bm->setROI1(validRoil);
	bm->setROI2(validRoir);
	bm->setPreFilterCap(31);
	bm->setBlockSize(SADWindowSize > 0 ? SADWindowSize : 9);
	bm->setMinDisparity(0);
	bm->setNumDisparities(numberOfDisparities);
	bm->setTextureThreshold(10);
	bm->setUniquenessRatio(15);
	bm->setSpeckleWindowSize(100);
	bm->setSpeckleRange(32);
	bm->setDisp12MaxDiff(1);

	//set sgbm
	sgbm->setPreFilterCap(63);
	int sgbmWinSize = SADWindowSize > 0 ? SADWindowSize : 3;
	sgbm->setBlockSize(sgbmWinSize);

	int cn = rectifyImageL.channels();

	sgbm->setP1(8*cn*sgbmWinSize*sgbmWinSize);
	sgbm->setP2(32*cn*sgbmWinSize*sgbmWinSize);
	sgbm->setMinDisparity(0);
	sgbm->setNumDisparities(numberOfDisparities);
	sgbm->setUniquenessRatio(10);
	sgbm->setSpeckleWindowSize(100);
	sgbm->setSpeckleRange(32);
	sgbm->setDisp12MaxDiff(1);

	if(alg==STEREO_HH)
	{
		sgbm->setMode(StereoSGBM::MODE_HH);
	}
	else if(alg==STEREO_SGBM)
	{
		sgbm->setMode(StereoSGBM::MODE_SGBM);
	}
	else if(alg==STEREO_3WAY)
	{
		sgbm->setMode(StereoSGBM::MODE_SGBM_3WAY);
	}

	int64 t = getTickCount();

	//compute left depth map
	if( alg == STEREO_BM )
	{
		bm->compute(rectifyImageL, rectifyImageR, leftDisp);
	}
	else if( alg == STEREO_SGBM || alg == STEREO_HH || alg == STEREO_3WAY )
	{
		sgbm->compute(rectifyImageL, rectifyImageR, leftDisp);
	}

	t = getTickCount() - t;
    printf("Time elapsed: %fms\n", t*1000/getTickFrequency());
}

void stereoCalculateDepthMap(const Mat &leftGray, const Mat &rightGray,
			Mat &cameraMatrixLeft, Mat &cameraMatrixRight, Mat &discoeffLeft, Mat &discoeffRight, Mat &R, Mat &T, Mat &depthMap, Mat &xyz)
{
	Mat rectifyImageL, rectifyImageR, q;
	Rect validRoil, validRoir;
	correctStereoImages(leftGray, rightGray, cameraMatrixLeft, cameraMatrixRight, discoeffLeft, discoeffRight, R, T,
									rectifyImageL, rectifyImageR, validRoil, validRoir, q);

	//get depth map
	stereoMatch(validRoil, validRoir, rectifyImageL, rectifyImageR, depthMap);

	//depth map to 3d distance
	reprojectImageTo3D(depthMap, xyz, q, true);
	xyz = xyz * 16 / 1000;//unint:m
}
